#include "mt3d.h"
#include "../utilities/matvec.h"

void mt3d::post_process()
{
    std::clog << "Processing MT stations ... \n";

    // Sort out MT stations
    mt_station tmp_station;
    for (size_t i = 0; i < 2; i++)
    {
        for (size_t j = 0; j < 2; j++)
        {
            tmp_station.impedance[i][j] = _zero;
            tmp_station.phi[i][j] = 0.0;
        }
        
        tmp_station.VMTF[i] = _zero;
    }

    for (size_t i = 0; i < node_num_; i++)
    {
        if (node_markers_[i] == station_marker_)
        {
            tmp_station.vert = mesh_nodes_.get(i);
            obs_site_.push_back(tmp_station);
        }
    }

    obs_num_ = obs_site_.size();

    // Find tetrahedrons that include mt stations
    for (size_t i = 0; i < elem_num_; i++)
    {
        if (elem_markers_[i] != air_marker_) // exclude the air layer to obtain smoother results
        {
            for (size_t j = 0; j < 4; j++)
            {
                if (node_markers_[mesh_elems_[i].vert[j]->id] == station_marker_)
                {
                    for (size_t k = 0; k < obs_num_; k++)
                    {
                        if (obs_site_[k].vert == mesh_elems_[i].vert[j])
                        {
                            obs_site_[k].host.push_back(mesh_elems_.get(i));
                            break;
                        }
                    }
                }
            }
        }
    }

    // Calculate observations
    gctl::array<double> total_vol(obs_num_, 0.0);
    for (size_t i = 0; i < obs_num_; i++)
    {
        for (size_t j = 0; j < obs_site_[i].host.size(); j++)
        {
            total_vol[i] += obs_site_[i].host[j]->volume();
        }
    }

    efem_sf::ordertype_e o_type;
    gctl::vertex3dc* tmp_vert[4];
    double ux, uy, uz, vx, vy, vz, dummy, tmp_vol;
    complex_d te_e[2], tm_e[2], te_h[3], tm_h[3];
    complex_d TE_E[2], TM_E[2], TE_H[3], TM_H[3];
    complex_d Emat[2][2], Hmat[2][2], Hmat_inv[2][2], Zmat[2][2];
    complex_d Tvec[2], Hvec[2], Hmat_det;
    double Zreal[2][2], Zreal_inv[2][2], Zimag[2][2], Phi[2][2];
    double Zreal_det;
    
    //complex_d te_ex, te_hy, tm_ey, tm_hx;
    //gctl::array<complex_d> TE_EX(obs_num_, _zero), TE_HY(obs_num_, _zero), TM_EY(obs_num_, _zero), TM_HX(obs_num_, _zero);

    for (size_t i = 0; i < obs_num_; i++)
    {
        TE_E[0] = TE_E[1] = TM_E[0] = TM_E[1] = _zero;
        TE_H[0] = TE_H[1] = TE_H[2] = TM_H[0] = TM_H[1] = TM_H[2] = _zero;
        for (size_t j = 0; j < obs_site_[i].host.size(); j++)
        {
            tmp_vert[0] = obs_site_[i].host[j]->vert[0];
            tmp_vert[1] = obs_site_[i].host[j]->vert[1];
            tmp_vert[2] = obs_site_[i].host[j]->vert[2];
            tmp_vert[3] = obs_site_[i].host[j]->vert[3];

            te_e[0] = te_e[1] = tm_e[0] = tm_e[1] = _zero;
            te_h[0] = te_h[1] = te_h[2] = tm_h[0] = tm_h[1] = tm_h[2] = _zero;
            //te_ex = te_hy = tm_ey = tm_hx = _zero;
            for (size_t d = 0; d < 6; d++)
            {
                if (elem_edge_reversed_[obs_site_[i].host[j]->id][d]) o_type = efem_sf::Reverse;
                else o_type = efem_sf::Normal;

                elsf_.tetrahedron(obs_site_[i].vert->x, obs_site_[i].vert->y, obs_site_[i].vert->z, 
                    tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, 
                    d, efem_sf::Value, o_type, ux, uy, uz, dummy, dummy, dummy);

                elsf_.tetrahedron(obs_site_[i].vert->x, obs_site_[i].vert->y, obs_site_[i].vert->z, 
                    tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, 
                    d, efem_sf::Curl, o_type, vx, vy, vz, dummy, dummy, dummy);

                // Note that we flipped the coordinate system
                te_e[1] += ux*X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                te_e[0] += uy*X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                tm_e[1] += ux*X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                tm_e[0] += uy*X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]];

                te_h[1] += vx*X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                te_h[0] += vy*X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                te_h[2] += -1.0*vz*X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                tm_h[1] += vx*X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                tm_h[0] += vy*X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                tm_h[2] += -1.0*vz*X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]];

                //te_ex += uy*X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                //te_hy += vx*X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]];

                //tm_ey += ux*X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
                //tm_hx += vy*X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]];
            }

            tmp_vol = obs_site_[i].host[j]->volume();
            TE_E[0] += te_e[0]*tmp_vol/total_vol[i];
            TE_E[1] += te_e[1]*tmp_vol/total_vol[i];
            TM_E[0] += tm_e[0]*tmp_vol/total_vol[i];
            TM_E[1] += tm_e[1]*tmp_vol/total_vol[i];

            TE_H[0] += te_h[0]*tmp_vol/total_vol[i];
            TE_H[1] += te_h[1]*tmp_vol/total_vol[i];
            TE_H[2] += te_h[2]*tmp_vol/total_vol[i];
            TM_H[0] += tm_h[0]*tmp_vol/total_vol[i];
            TM_H[1] += tm_h[1]*tmp_vol/total_vol[i];
            TM_H[2] += tm_h[2]*tmp_vol/total_vol[i];

            //TE_EX[i] += te_ex*tmp_vol/total_vol[i];
            //TE_HY[i] += te_hy*tmp_vol/total_vol[i];
            //TM_EY[i] += tm_ey*tmp_vol/total_vol[i];
            //TM_HX[i] += tm_hx*tmp_vol/total_vol[i];
        }

        TE_H[0] /= _1i*omega_*_mu0;
        TE_H[1] /= _1i*omega_*_mu0;
        TE_H[2] /= _1i*omega_*_mu0;
        TM_H[0] /= _1i*omega_*_mu0;
        TM_H[1] /= _1i*omega_*_mu0;
        TM_H[2] /= _1i*omega_*_mu0;

        //TE_HY[i] /= _1i*omega_*_mu0;
        //TM_HX[i] /= _1i*omega_*_mu0;

        le_te_.calculate_EH(-1.0*obs_site_[i].vert->z, 0.0); // Zero obique angle
        TE_E[0] += le_te_.get_Ex();
        TE_E[1] += le_te_.get_Ey();
        TE_H[0] += le_te_.get_Hx();
        TE_H[1] += le_te_.get_Hy();
        TE_H[2] += le_te_.get_Hz();

        //le_.calculate_EH(-1.0*obs_site_[i].vert->z, 0.0); // Zero obique angle
        //TE_E[0] += le_.get_Ex();
        //TE_E[1] += le_.get_Ey();
        //TE_H[0] += le_.get_Hx();
        //TE_H[1] += le_.get_Hy();
        //TE_H[2] += le_.get_Hz();

        //TE_EX[i] += le_te_.get_Ex();
        //TE_HY[i] += le_te_.get_Hy();

        le_tm_.calculate_EH(-1.0*obs_site_[i].vert->z, 0.0); // Zero obique angle
        TM_E[0] += le_tm_.get_Ex();
        TM_E[1] += le_tm_.get_Ey();
        TM_H[0] += le_tm_.get_Hx();
        TM_H[1] += le_tm_.get_Hy();
        TM_H[2] += le_tm_.get_Hz();

        //TM_E[0] -= le_.get_Ey();
        //TM_E[1] += le_.get_Ex();
        //TM_H[0] -= le_.get_Hy();
        //TM_H[1] += le_.get_Hx();
        //TM_H[2] += le_.get_Hz();

        //TM_EY[i] += le_tm_.get_Ey();
        //TM_HX[i] += le_tm_.get_Hx();

        //TE_E[0] += le_.get_Ey();
        //TE_E[1] += le_.get_Ex();
        //TM_E[0] += le_.get_Ey();
        //TM_E[1] += le_.get_Ex();

        //TE_H[0] += le_.get_Hy();
        //TE_H[1] += le_.get_Hx();
        //TE_H[2] += -1.0*le_.get_Hz();
        //TM_H[0] += le_.get_Hy();
        //TM_H[1] += le_.get_Hx();
        //TM_H[2] += -1.0*le_.get_Hz();

        Emat[0][0] = TE_E[0]; Emat[0][1] = TM_E[0];
        Emat[1][0] = TE_E[1]; Emat[1][1] = TM_E[1];

        Hmat[0][0] = TE_H[0]; Hmat[0][1] = TM_H[0];
        Hmat[1][0] = TE_H[1]; Hmat[1][1] = TM_H[1];
        Hvec[0] = TE_H[2];
        Hvec[1] = TM_H[2];

        // Calculate impedance
        INVERT_2X2(Hmat, Hmat_det, Hmat_inv);
        MAT_PRODUCT_2X2(Emat, Hmat_inv, Zmat);
        
        // Calculate VMTF
        Hmat[0][0] = TE_H[0]; Hmat[0][1] = TE_H[1];
        Hmat[1][0] = TM_H[0]; Hmat[1][1] = TM_H[1];

        INVERT_2X2(Hmat, Hmat_det, Hmat_inv);
        MAT_VEC_2X2(Hmat_inv, Hvec, Tvec);

        // Calculate phase tensor
        C2R_2X2(Zmat, Zreal, Zimag);
        INVERT_2X2(Zreal, Zreal_det, Zreal_inv);
        MAT_PRODUCT_2X2(Zreal_inv, Zimag, Phi);

        // Save results to site
        COPY_2X2(Zmat, obs_site_[i].impedance);
        COPY_2X2(Phi, obs_site_[i].phi);
        obs_site_[i].VMTF[0] = Tvec[0];
        obs_site_[i].VMTF[1] = Tvec[1];

        //obs_site_[i].Zxy = TE_EX[i]/TE_HY[i];
        //obs_site_[i].Zyx = TM_EY[i]/TM_HX[i];
    }

    return;
}
